import os
from optparse import OptionParser
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,pipeline

global g_EN_DATA
g_EN_DATA = "en_vaild.en"

global g_ZH_DATA
g_ZH_DATA = "zh_vaild.zh"

global g_EN_OUT_PUT
g_EN_OUT_PUT = "en_zh_result.zh"

global g_ZH_OUT_PUT
g_ZH_OUT_PUT = "zh_en_result.en"

# 最大翻译语句数
global MAX_TRA
MAX_TRA = 200

###################################
# load_model： 加载模型文件
###################################
def load_model():
    # 英文-中文
    # 载入序列化英文-中文模型
    tokenizer1 = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh",cache_dir='./en-zh-model')

    # 载入翻译英文-中文模型
    model1 = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-zh",cache_dir='./en-zh-model')

    # 载入模型
    translation1 = pipeline("translation_en_to_zh", model=model1, tokenizer=tokenizer1)

    # 中文-英文
    # 载入序列化中文-英文模型
    tokenizer2 = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en",cache_dir='./zh-en-model')

    # 载入翻译中文-英文模型
    model2 = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-zh-en",cache_dir='./zh-en-model')

    # 载入模型
    translation2 = pipeline("translation_zh_to_en", model=model2, tokenizer=tokenizer2)

    return [translation1,translation2];

###################################
# translation： 
###################################
def translation(translation1,translation2):

    f_en=open(g_EN_DATA,"r",encoding='UTF-8')
    f_en_output=open(g_EN_OUT_PUT,"w",encoding='UTF-8')

    f_zh=open(g_ZH_DATA,"r",encoding='UTF-8')
    f_zh_output=open(g_ZH_OUT_PUT,"w",encoding='UTF-8')

    text_ens=f_en.readlines()

    text_zhs=f_zh.readlines()

    trans_num_en_zh=min(text_ens.__len__(),MAX_TRA)
    trans_num_zh_en=min(text_zhs.__len__(),MAX_TRA)
    
    # 英文-中文
    print("英文翻译中文:%d条\n"%trans_num_en_zh)
    count1=0
    
    text_ens=text_ens[:trans_num_en_zh]
    for text_en in text_ens:
        # 去掉换行符
        text_en=text_en[:-2]
        translated_text1 = translation1(text_en, max_length=512)[0]['translation_text']
        f_en_output.write(translated_text1+"\n")
        print("第%d条翻译成功"%count1)
        count1+=1


    # 中文-英文
    print("中文翻译英文:%d条\n"%trans_num_zh_en)
    count2=0
    
    text_zhs=text_zhs[:trans_num_zh_en]
    for text_zh in text_zhs:
        # 去掉换行符
        text_zh=text_zh[:-2]
        translated_text2 = translation2(text_zh, max_length=512)[0]['translation_text']
        f_zh_output.write(translated_text2+"\n")
        print("第%d条翻译成功"%count2)
        count2+=1

    f_zh.close()
    f_zh_output.close()

    return None;


###################################
# main： 主逻辑
###################################
def main():
    print("main start.....")

    model=load_model()
    translation(model[0],model[1])

    print("main end...")

if __name__ == '__main__':
	main()

















